# -*- coding: utf-8 -*-
# 作者：秋禊
import numpy as np
from PIL import Image
import os
import cv2

# 人脸数据路径
path = 'Facedata'

# 创建识别模型，使用LBPHFace算法识别，Confidence评分低于50是可靠
# LBP是一种特征提取方式，能提取出图像的局部的纹理特征
recognizer = cv2.face.LBPHFaceRecognizer_create()
detector = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")

def get_images_and_labels(path):
    # 所有图片路径
    image_paths = [os.path.join(path, f) for f in os.listdir(path)]  # join函数的作用？

    faceSamples = []
    ids = []
    for image_path in image_paths:
        # 转换为灰度
        pil_img = Image.open(image_path).convert('L')
        # PIL_img = cv2.cvtColor(imagePath, cv2.COLOR_BGR2GRAY)
        # 把灰度图片变为矩阵
        img_numpy = np.array(pil_img, 'uint8')
        # 把图片路径，以 。为界分割，第一项为user的名字
        id = int(os.path.split(image_path)[-1].split(".")[1])

        # 把检测到的对象图像大小的值，作为矩阵返回
        faces = detector.detectMultiScale(img_numpy)
        # 矩阵的每一行的四个元素，对应着 x, y, w, h
        for (x, y, w, h) in faces:
            faceSamples.append(img_numpy[y: y+h, x: x+w])
            ids.append(id)
    return faceSamples, ids


print('人脸模型的训练需要一段时间，请稍候 ...')
faces, ids = get_images_and_labels(path)
recognizer.train(faces, np.array(ids))

recognizer.write(r'face_trainer\trainer.yml')
print("{0} 个面部。".format(len(np.unique(ids))))
